#!/usr/bin/python

"""
This program draws genome comparison plot between two given genomes. It does similar things to
Artemis-based genome comparison at DNA level but instead of using DNA sequence, the program
uses reciprocal orthologous BLAST hits to chart where the orthologs between two given genomes
are located. This helps to visualize if any syntenous regions for a certain class of proteins
are present and also helps to se genome rearrangements.

Usage: 
Examples: 
Run in this folder: /host/Users/JS/UH-work/gloeobacter/final_work/comparisons/orthologs/orthomcl
dissertation_ReciprocalBestHitPlot.py ../../../annotation/GKIL.v6.gbf ../../NC_005125.1.gbk pair_GKIL_58011.txt cogs.t.list
dissertation_ReciprocalBestHitPlot.py ../cyano/61581.refseq.gbk ../cyano/61607.refseq.gbk pair_61581_61607.txt cogs.t.list
"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC

##Regex
gkil = re.compile('GKIL\|(\w+)')
gvio = re.compile('58011\|(\w+)') # 2nd genome to compare with. can change to see different plots

#gkil = re.compile('61581\|(\w+)')
#gvio = re.compile('61607\|(\w+)')


cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')

#hex colors

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#D2B9D3',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
    '-' : '#EDEDED'
}

##Genome one Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    if feat.type == 'CDS':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat


##Genome two Genbank file
g2seq = SeqIO.read(sys.argv[2], "gb")
g2length = len(g2seq.seq)
g2featdict = {}
for feat in g2seq.features:
    if feat.type == 'CDS':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat


larger_genome = 0

if g1length > g2length:
    larger_genome = g1length
else:
    larger_genome = g2length

pairfile = open(sys.argv[3], "rU")
pfl = pairfile.readlines()

totalg1hits = 0
totalg2hits = 0

cogcatfile = open(sys.argv[4], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

"""
parse orthologs in pairs file. Pair file should contain filtered orthologous
groups that should look like this:
CYANO15211: 57767|Npun_F1941 58043|Ava_0214 57803|alr2398 GKIL|GKIL_3674 59025|PCC7424_4452 58011|glr3849
CYANO15212: 57767|Npun_R1991 GKIL|GKIL_3611 58167|AM1_4177 57767|Npun_F2216 59435|Cyan7425_3596 58011|glr1283
CYANO15213: 58011|glr0952 59435|Cyan7425_0962 GKIL|GKIL_0709 57803|alr7163 58043|Ava_B0061 58167|AM1_2965
CYANO15215: 57767|Npun_R2251 GKIL|GKIL_4327 58011|glr1730 57767|Npun_R5324 58043|Ava_2173 57925|Tery_0799
"""

glist = []

for line in pfl:
    t = line.split(' ')
    oname = t[0]
    g1x = 0 #x-coord for genome1
    g2x = 0 #x-coord for genome2
    #cogcolor = "#385E0F" #if no COGs, this is the base color
    #cogcolor = "#EEE0E5"
    cogcolor = "#EDEDED"
    for i in t[1:]: #note that each line may contain multiple hits. eg: 2 vs 3
        tmp = i.strip()
        if gkil.match(tmp):
            pat1 = gkil.match(tmp)
            ltag = pat1.group(1)
            if ltag in g1featdict:
                g1feat = g1featdict[ltag]
                g1start = g1feat.location._start.position
                g1stop = g1feat.location._end.position
                g1mid = g1start + ((g1stop - g1start)/2.0) #places the point to middle of gene
                g1x = g1mid
                g1desc = g1feat.qualifiers['product'][0]
                cog = ""
                if g1feat.qualifiers.has_key('note'):
                    cog = g1feat.qualifiers['note'][0]
                    if cog in cogcatdict:
                        cogcolor = cogdict[cogcatdict[cog]]
                #print '{0:10}\t{1}\t{2}\t{3}'.format(ltag, g1start, g1stop, g1desc)
                totalg1hits += 1
        elif gvio.match(tmp):
            pat2 = gvio.match(tmp)
            ltag = pat2.group(1)
            if ltag in g2featdict:
                g2feat = g2featdict[ltag]
                g2start = g2feat.location._start.position
                g2stop = g2feat.location._end.position
                g2mid = g2start + ((g2stop - g2start)/2.0)
                g2x = g2mid
                g2desc = g2feat.qualifiers['product'][0]
                #print '{0:10}\t{1}\t{2}\t{3}'.format(ltag, g2start, g2stop, g2desc)
                totalg2hits += 1
    glist.append((g1x, g2x, cogcolor)) #append the x-coordinates for genome 1 and 2, and cog color

##For lines representing the genomes. Needed to set upper and lower bounds above and below x-y segments
genome1x = [2, g1length]
genome1y = [2, 2]
genome2x = [2, g2length]
genome2y = [10, 10]

"""
#Function to flip the genome coordinate if needed. Keep for further testing.
flippedlist = []
glistlen = len(glist)
x = 0
beginning = 0
end = -1
while x < glistlen:
    x1 = glist[beginning][0]
    x2 = glist[end][1]
    cc = glist[beginning][2]
    beginning = beginning + 1
    end = end - 1
    flippedlist.append((x1, x2, cc)) #appending flipped list
    x += 1
"""

##Start plotting
#fig = plt.figure(1, figsize=(15,8)) #to be used with 211
fig = plt.figure(1, figsize=(15,4))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.axis([0, larger_genome, 0, 14])


for i in glist:
    x = [i[0], i[1]] #i[0] is genome 1 coord and i[1] is genome 2 coord
    y = [3, 9] #y axes
    c = i[2]
    #ax1.plot(x, y, "o-", alpha=0.2)
    #ax1.plot(x, y, "s-", alpha=0.2)
    #ax1.plot(x, y, color=c, marker="o", alpha=0.2)
    ax1.plot(x, y, color=c, marker="s", alpha=0.2)
    #ax1.plot(x, y, "#333366", alpha=0.2) #should set the color based on COG categories


"""
#to plot flipped order
for i in flippedlist:
    x = [i[0], i[1]] #i[0] is genome 1 coord and i[1] is genome 2 coord
    y = [3, 9] #y axes
    c = i[2]
    #ax1.plot(x, y, "o-", alpha=0.2)
    #ax1.plot(x, y, "s-", alpha=0.2)
    ax1.plot(x, y, color=c, marker="o", alpha=0.2)
"""

ax1.annotate('GKIL', xy=(4850000, 3), horizontalalignment='center', verticalalignment='center', fontsize=10)
ax1.annotate('GVIO', xy=(4850000, 9), horizontalalignment='center', verticalalignment='center', fontsize=10)

#ax1.annotate('WH 8102', xy=(2400000, 2.5), horizontalalignment='center', verticalalignment='center', fontsize=10)
#ax1.annotate('WH 7803', xy=(2400000, 9.5), horizontalalignment='center', verticalalignment='center', fontsize=10)



#rect = Rectangle((500000, 2.5), 20000, 0.25, facecolor="#2B60DE", alpha=0.5)
#plt.gca().add_patch(rect)
#bbox = Bbox.from_bounds(500000, 2.5, 20000, 0.25)


#Draw legend box for COG categories

coglist = []
for k, v in cogdict.iteritems():
    coglist.append((k,v))

coglist.sort()

ccounts = len(coglist)
a = 0
xstart = 2500000
#increment = 20000 #for synecoccus
increment = 40000 #for gloeobacter
while a < ccounts:
    fc = coglist[a][1]
    tx = coglist[a][0]
    #rect = Rectangle((xstart, 2.5), 20000, 0.25, facecolor=fc, alpha=0.5) #for synecococcus
    rect = Rectangle((xstart, 2.5), 40000, 0.25, facecolor=fc, alpha=0.5) #for gloeobacter
    plt.gca().add_patch(rect)
    ax1.annotate(tx, xy=(xstart+20000, 2.3), horizontalalignment='center', verticalalignment='center', fontsize=8)
    a += 1
    xstart = xstart + increment
ax1.annotate('COG categories', xy=(2000000, 2.3), fontsize=8)

frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)
plt.show()

#print totalg1hits



##Start plotting
#fig = plt.figure(1, figsize=(15,8))

#xlength = larger_genome + 1000

#ax1.axis([0, xlength, 0, 14])

pairfile.close()
#gl1.close()
#gl2.close()

